/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <cstring>
#include <iostream>
#include "utils.h"
#include "acl_nn_operation.h"

namespace atb_plugin {

AclNNOperation::AclNNOperation(const std::string &opName) : opName_(opName)
{
}

AclNNOperation::~AclNNOperation()
{
    std::cout << "AclNNOperation deconstructor" << std::endl;
    this->DestoryOperation();
}

std::string AclNNOperation::GetName() const { return this->opName_; }

void AclNNOperation::DestoryOperation()
{
    if (this->aclExecutor == nullptr) {
        return;
    }

    this->aclExecutor = nullptr;

    // 清空aclTensor
    for (size_t i = 0; i < this->aclInTensors.size(); ++i) {
        aclDestroyTensor(this->aclInTensors[i]->tensor);
    }
    this->aclInTensors.clear();

    for (size_t i = 0; i < this->aclOutTensors.size(); ++i) {
        aclDestroyTensor(this->aclOutTensors[i]->tensor);
    }
    this->aclOutTensors.clear();
}

atb::Status AclNNOperation::Setup(const atb::VariantPack &variantPack, uint64_t &workspaceSize, atb::Context *context)
{
    std::cout << this->opName_ << "setup start" << std::endl;
    if (context == nullptr) {
        std::cout << this->opName_ << "setup context is null" << std::endl;
        return atb::ERROR_INVALID_PARAM;
    }

    int ret = CreateAclNNOpCache(variantPack);
    if (ret != 0) {
        std::cout << this->opName_ << "call CreateAclNNOpCache fail" << std::endl;
        return ret;
    }

    ret = SetAclNNWorkspaceExecutor();
    if (ret != 0) {
        std::cout << this->opName_ << "call SetAclNNWorkspaceExecutor fail" << std::endl;
        return ret;
    }

    workspaceSize = this->workspaceSize;
    
    std::cout << "setup end. " << "WorkspaceSize:" << workspaceSize << std::endl;
    return atb::NO_ERROR;
}

atb::Status AclNNOperation::CreateAclNNOpCache(const atb::VariantPack &variantPack)
{
    int ret = CreateAclNNVariantPack(variantPack);
    if (ret != 0) {
        std::cout << this->opName_ << "call CreateAclNNVariantPack fail" << std::endl;
        return atb::ERROR_CANN_ERROR;
    }
    return atb::NO_ERROR;
}

atb::Status AclNNOperation::Execute(const atb::VariantPack &variantPack, uint8_t *workspace, uint64_t workspaceSize,
                                    atb::Context *context)
{
    std::cout << this->opName_ << "execute start" << std::endl;
    if (!context) {
        std::cout << this->opName_ << "execute fail, context param is null" << std::endl;
        return atb::ERROR_INVALID_PARAM;
    }

    aclrtStream stream = context->GetExecuteStream();
    if (!stream) {
        std::cout << this->opName_ << "execute fail, execute stream in context is null" << std::endl;
        return atb::ERROR_INVALID_PARAM;
    }

    int ret = ExecuteAclNNOp(workspace, stream);
    if (ret != 0) {
        std::cout << this->opName_ << "call ExecuteAclNNOp fail" << std::endl;
        return atb::ERROR_CANN_ERROR;
    }

    std::cout << this->opName_ << "execute end" << std::endl;
    return atb::NO_ERROR;
}

} // namespace atb_plugin
